//
//  KImageSegmentator.m
//  ToneMapper
//
//  Created by Andrey on 27/03/2009.
//  Copyright 2009 Karma Software. All rights reserved.
//

#import "KImageSegmentator.h"
#import "KImage.h"
#import "KSimpleUIntArray.h"
#import "KSimpleFloatArray.h"
#import "KSegmentationMap.h"
#import "KPixel.h"

#define HUE_WEIGHT 1.0
#define DISTANCE_WEIGHT 0.0

typedef struct _KSegment 
{
	CGFloat meanValue;
	NSUInteger pixelCount;
	CGFloat centroid;
} KSegment, *KSegmentRef, *KSegmentArray;

typedef struct _K2DSegment 
{
	NSPoint meanValue;
	NSUInteger pixelCount;
	NSPoint centroid;
} K2DSegment, *K2DSegmentArray;

CGFloat GetEuclideanDistanceBetweenPoints(NSPoint p1, NSPoint p2)
{
	CGFloat deltaX = p2.x - p1.x;
	CGFloat deltaY = p2.y - p1.y;
	
	return (CGFloat)sqrt(deltaX * deltaX + deltaY * deltaY);
}

NSUInteger GetMinElementIndexInFloatArray(CGFloat* array, NSUInteger elementCount)
{
	CGFloat minValue = 1e30;
	NSUInteger minValueIndex = 0;
	
	for (int i = 0; i < elementCount; i++)
	{
		if (array[i] < minValue)
		{
			minValue = array[i];
			minValueIndex = i;
		}
	}
	
	return minValueIndex + 1;
}

@implementation KImageSegmentator

+ (KSegmentationMap*) segmentImage: (KImage*)image 
	   usingKMeansWithSegmentCount: (NSUInteger)segmentCount 
				   andUsingChannel: (enum KImageChannel)channel
{
	// Randomly choose initial segment centroids
	KSegmentArray segments = (KSegmentArray)malloc(sizeof(KSegment) * segmentCount);
	bzero(segments, sizeof(KSegment) * segmentCount);
	for (int i = 0; i < segmentCount; i++)
	{
		segments[i].centroid = (CGFloat)random() / (CGFloat)RAND_MAX;
//		segments[i].centroid = (CGFloat)i / (segmentCount - 1);
	}
	
	// Distances in hue space from current pixel to each of centroids
	CGFloat* distancesToCentroids = (CGFloat*)malloc(sizeof(CGFloat) * segmentCount);

	// Initial segmentation map
	NSUInteger imagePixelCount = image.width * image.height;
	KSimpleUIntArray* segmentationMap = [KSimpleUIntArray arrayWithCapacity:imagePixelCount];
	
	// Segmentation process
	Boolean segmentationDidChange = YES;
	while (segmentationDidChange) 
	{
		// Restore segment states
		segmentationDidChange = NO;
		NSLog(@"---------- New segmentation pass ------------");
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].meanValue = 0.0;
			segments[i].pixelCount = 0;
			NSLog(@"Segment #%i centroid:\t %.4f", i, segments[i].centroid);
		}
		
		// Compute new segmentation
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			KPixel* currentPixel = [image pixelAtX:x y:y];
				
			CGFloat channelValueAtCurrentPixel = 
			[currentPixel redundantColorVector].data[channel];
			
			for (int i = 0; i < segmentCount; i++) 
			{
				distancesToCentroids[i] = 
				fabs(channelValueAtCurrentPixel - segments[i].centroid);
			}
			
			NSUInteger closestCentroid = 
			GetMinElementIndexInFloatArray(distancesToCentroids, segmentCount);
			
			if (closestCentroid != segmentationMap.data[index])
				segmentationDidChange = YES;
			
			segmentationMap.data[index] = closestCentroid;
			
			segments[closestCentroid - 1].pixelCount++;
			segments[closestCentroid - 1].meanValue += channelValueAtCurrentPixel;
		}
		
		// Compute new centroids
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].centroid = segments[i].meanValue / segments[i].pixelCount;
		}
	}
	
	KSimpleFloatArray* distancesBetweenCentroids = 
	[KSimpleFloatArray arrayWithCapacity:(segmentCount - 1)];
	
	for (int i = 1; i < segmentCount; i++)
	{
		distancesBetweenCentroids.data[i - 1] = 
		fabs(segments[i].centroid - segments[i - 1].centroid);
	}
	
	CGFloat variance = [distancesBetweenCentroids maxValue];
	
	KSegmentationMap* result = 
	[[KSegmentationMap alloc] 
	 initWithPixelCount:imagePixelCount andSegmentCount:segmentCount];
	
	for (NSUInteger i = 0; i < segmentCount; i++)
	{
		CGFloat centroid = segments[i].centroid;
		
		// Compute probability map
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			KPixel* currentPixel = [image pixelAtX:x y:y];
			
			CGFloat channelValueAtCurrentPixel = 
			[currentPixel redundantColorVector].data[channel];
				
				CGFloat distanceToCentroid = centroid - channelValueAtCurrentPixel;
				CGFloat distanceToCentroidSquared = 
				distanceToCentroid * distanceToCentroid;
				
				result.segments[i].data[index] = 
				exp(-distanceToCentroidSquared / (2 * variance * variance));
		}
											 
		
		result.centroids.data[i] = centroid;
	}
	
	// Clean up
	free(distancesToCentroids);
	free(segments);
	
	return result;
}

+ (K2DSegmentationMap*) segmentImageByChromaticity: (KImage*)image 
					   usingKMeansWithSegmentCount: (NSUInteger)segmentCount
{
	NSUInteger imagePixelCount = image.width * image.height;
	
	K2DSegmentArray segments = (K2DSegmentArray)malloc(sizeof(K2DSegment) * segmentCount);
	
	// Initial segments in CIEa*b* space (!) despite x and y
	for (int i = 0; i < segmentCount; i++)
	{

		segments[i].centroid.x = (CGFloat)random() / (CGFloat)RAND_MAX;
		segments[i].centroid.y = (CGFloat)random() / (CGFloat)RAND_MAX;

		/*
		segments[i].centroid.x = (CGFloat)i / segmentCount;
		segments[i].centroid.y = (CGFloat)i / segmentCount;
		 */
	}
	
	// Distances to centroids
	KSimpleFloatArray* distancesToCentroids = 
	[[KSimpleFloatArray alloc] initWithCapacity:segmentCount];
	
	// Simple segmentation map
	KSimpleUIntArray* simpleSegmentationMap =
	[[KSimpleUIntArray alloc] initWithCapacity:imagePixelCount];
	
	Boolean segmentationDidChange = YES;
	while (segmentationDidChange) 
	{
		// Restore segment states
//		NSLog(@"---------- New segmentation pass ------------");
		segmentationDidChange = NO;
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].meanValue = NSMakePoint(0.0, 0.0);
			segments[i].pixelCount = 0;
/*			NSLog(@"Segment #%i centroid:\t {%.4f, %.4f}", i, 
				  segments[i].centroid.x, 
				  segments[i].centroid.y);
*/		}
		
		// Compute new segmentation
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			
			KPixel* currentPixel = [image pixelAtX:x y:y];
			
			NSPoint abValuesAtCurrentPixel = 
			NSMakePoint([currentPixel a], [currentPixel b]);
			
			for (int i = 0; i < distancesToCentroids.count; i++)
			{
				distancesToCentroids.data[i] = 
				GetEuclideanDistanceBetweenPoints(abValuesAtCurrentPixel, 
												  segments[i].centroid);
			}
			
			NSUInteger closestCentroid = [distancesToCentroids minValueIndex];
			assert(closestCentroid != -1);
			if (closestCentroid != simpleSegmentationMap.data[index])
				segmentationDidChange = YES;
			
			simpleSegmentationMap.data[index] = closestCentroid;
			
			NSPoint previousMeanValue = segments[closestCentroid].meanValue;
			segments[closestCentroid].meanValue =
			NSMakePoint(previousMeanValue.x + abValuesAtCurrentPixel.x, 
						previousMeanValue.y + abValuesAtCurrentPixel.y);
			segments[closestCentroid].pixelCount++;
		}
		
		// Compute new centroids
		
		for (int i = 0; i < segmentCount; i++)
		{
			segments[i].centroid =
			NSMakePoint(segments[i].meanValue.x / segments[i].pixelCount,
						segments[i].meanValue.y / segments[i].pixelCount);
		}
	}
	
	KSimpleFloatArray* distancesBetweenCentroids = 
	[KSimpleFloatArray arrayWithCapacity:(segmentCount - 1)];
	
	for (int i = 1; i < segmentCount; i++)
	{
		distancesBetweenCentroids.data[i - 1] = 
		GetEuclideanDistanceBetweenPoints(segments[i].centroid, segments[i - 1].centroid);
	}
	
	CGFloat variance = [distancesBetweenCentroids maxValue];
	
	// Form final segmentation map 
	K2DSegmentationMap* result = 
	[[K2DSegmentationMap alloc] 
	 initWithPixelCount:imagePixelCount andSegmentCount:segmentCount];
	
	for (NSUInteger i = 0; i < segmentCount; i++)
	{
		NSPoint centroid = segments[i].centroid;
		
		// Compute probability map
		ForEachImagePixel(image)
		{
			NSUInteger index = y * image.width + x;
			
			KPixel* currentPixel = [image pixelAtX:x y:y];
			
			NSPoint abValuesAtCurrentPixel = 
			NSMakePoint([currentPixel a], [currentPixel b]);
			
			CGFloat distanceToCentroid = 
			GetEuclideanDistanceBetweenPoints(abValuesAtCurrentPixel, 
											  centroid);
			
			CGFloat distanceToCentroidSquared = 
			distanceToCentroid * distanceToCentroid;
			
			result.segments[i].data[index] = 
			exp(-distanceToCentroidSquared / (2 * variance * variance));
		}
		
		result.centroidPoints[i] = centroid;
	}
	
	[distancesToCentroids release];
	free(segments);
	
	return result;
}

@end
